﻿#include "VecAdd.h"


VecAdd::VecAdd(ID3D12Device* device, ID3D12GraphicsCommandList* cmdList,int NumDataElements)
{
    mNumDataElements = NumDataElements;
    
    std::vector<Data> dataA(mNumDataElements);
    std::vector<Data> dataB(mNumDataElements);

    for(int i = 0; i < mNumDataElements; ++i)
    {
        dataA[i].v1 = XMFLOAT3((float)i, (float)i, (float)i);
        dataA[i].v2 = XMFLOAT2((float)i, 0);

        dataB[i].v1 = XMFLOAT3((float)-i, (float)i, 0.0f);
        dataB[i].v2 = XMFLOAT2(0.0f, (float)-i);
    }

    UINT64 byteSize = dataA.size() * sizeof(Data);
    
    mInputBufferA = d3dUtil::CreateDefaultBuffer(
        device,
        cmdList,
        dataA.data(),
        byteSize,
        mInputUploadBufferA
        );

    mInputBufferB = d3dUtil::CreateDefaultBuffer(
        device,
        cmdList,
        dataB.data(),
        byteSize,
        mBufferUpLoadB
        );

    auto outHeapPro = CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_DEFAULT);
    auto outBuffDes = CD3DX12_RESOURCE_DESC::Buffer(byteSize, D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS);
    ThrowIfFailed(device->CreateCommittedResource(
        &outHeapPro,
        D3D12_HEAP_FLAG_NONE,
        &outBuffDes,
        D3D12_RESOURCE_STATE_UNORDERED_ACCESS,
        nullptr,
        IID_PPV_ARGS(&mOutputBuffer)
        ));

    auto readBackHeapPro = CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_READBACK);
    auto readBackDes = CD3DX12_RESOURCE_DESC::Buffer(byteSize);
    ThrowIfFailed(device->CreateCommittedResource(
        &readBackHeapPro,
        D3D12_HEAP_FLAG_NONE,
        &readBackDes,
        D3D12_RESOURCE_STATE_COPY_DEST,
        nullptr,
        IID_PPV_ARGS(&mReadBackBuffer)));
    
}

void VecAdd::DoComputeWork(
    ID3D12CommandAllocator* allocator,
    ID3D12GraphicsCommandList* cmdList,
    ID3D12RootSignature* rootSig,
    ID3D12PipelineState* pso)
{
    ThrowIfFailed(allocator->Reset());
    ThrowIfFailed(cmdList->Reset(allocator,pso));
    
    cmdList->SetPipelineState(pso);
    cmdList->SetComputeRootSignature(rootSig);

    cmdList->SetComputeRootShaderResourceView(0,mInputBufferA->GetGPUVirtualAddress());
    cmdList->SetComputeRootShaderResourceView(1,mInputBufferB->GetGPUVirtualAddress());
    cmdList->SetComputeRootUnorderedAccessView(2,mOutputBuffer->GetGPUVirtualAddress());

    cmdList->Dispatch(1,1,1);

    auto transition = CD3DX12_RESOURCE_BARRIER::Transition(mOutputBuffer.Get(),
        D3D12_RESOURCE_STATE_UNORDERED_ACCESS,
        D3D12_RESOURCE_STATE_COPY_SOURCE
        );
    cmdList->ResourceBarrier(1, &transition);
    cmdList->CopyResource(mReadBackBuffer.Get(),mOutputBuffer.Get());

    auto transition2 = CD3DX12_RESOURCE_BARRIER::Transition(mOutputBuffer.Get(),
        D3D12_RESOURCE_STATE_COPY_SOURCE,
        D3D12_RESOURCE_STATE_COMMON
        );
    cmdList->ResourceBarrier(1,&transition2);
}

void VecAdd::AfterComputeWork()
{    
    // Map the data so we can read it on CPU.
    Data* mappedData = nullptr;
    ThrowIfFailed(mReadBackBuffer->Map(0, nullptr, reinterpret_cast<void**>(&mappedData)));

    std::ofstream fout(L"..\\..\\Output\\CSoutput\\results.txt");

    for(int i = 0; i < mNumDataElements; ++i)
    {
        fout << "(" << mappedData[i].v1.x << ", " << mappedData[i].v1.y << ", " << mappedData[i].v1.z <<
            ", " << mappedData[i].v2.x << ", " << mappedData[i].v2.y << ")" << std::endl;
    }

    mReadBackBuffer->Unmap(0,nullptr);
}
